import os
import json
import numpy as np
import cv2
import math


# 加载 JSON 文件
def load_json(json_path):
    with open(json_path, 'r') as f:
        data = json.load(f)
    return data


# 提取用于计算外展角和前倾角的关键点
def extract_keypoints(data):
    # 使用点 1~6 进行右侧拟合，8~13 进行左侧拟合
    right_labels = ['Rlong1', 'Rlong2', 'Rass1', 'Rass2', 'Rass3', 'Rass4']
    left_labels = ['Llong1', 'Llong2', 'Lass1', 'Lass2', 'Lass3', 'Lass4']

    right_points = []
    left_points = []

    for shape in data['shapes']:
        label = shape['label']
        point = shape['points'][0]  # 提取点坐标
        if label in right_labels:
            right_points.append(tuple(point))
        elif label in left_labels:
            left_points.append(tuple(point))

    return np.array(right_points, dtype=np.float32), np.array(left_points, dtype=np.float32)


# 计算向量夹角
def calculate_angle(v1, v2):
    v1 = np.array(v1)
    v2 = np.array(v2)

    # 防止除以零错误
    norm_v1 = np.linalg.norm(v1)
    norm_v2 = np.linalg.norm(v2)

    if norm_v1 == 0 or norm_v2 == 0:
        print("Warning: One of the vectors has zero length, skipping this angle calculation.")
        return None  # 返回 None，表示无法计算

    dot_product = np.dot(v1, v2)

    # 防止浮点误差导致 dot_product 超过 [-1, 1] 范围
    cos_theta = dot_product / (norm_v1 * norm_v2)
    cos_theta = np.clip(cos_theta, -1.0, 1.0)  # 将值限制在 [-1, 1] 范围内

    angle_rad = np.arccos(cos_theta)

    return angle_rad  # 返回弧度


# 拟合椭圆并获取最大直径D和其他椭圆信息
def fit_ellipse_and_get_major_axis(points, json_path, side):
    if len(points) < 5:  # 椭圆拟合需要至少5个点
        print(
            f"Error: Not enough points to fit ellipse for {side} side in file {json_path}. Need at least 5 points, but got {len(points)}.")
        return None, None, None

    ellipse = cv2.fitEllipse(np.array(points))
    center, axes, angle = ellipse
    major_axis_length = max(axes)  # 获取椭圆的最大直径 D
    return major_axis_length, center, ellipse


# 计算从点 M 到椭圆轮廓的垂直距离 p
def calculate_p(ellipse, M):
    center, axes, angle = ellipse
    # 将 M 转换到椭圆的局部坐标系中
    cos_angle = np.cos(np.radians(-angle))
    sin_angle = np.sin(np.radians(-angle))

    # 计算局部坐标中的 M 点
    local_M_x = (M[0] - center[0]) * cos_angle + (M[1] - center[1]) * sin_angle
    local_M_y = -(M[0] - center[0]) * sin_angle + (M[1] - center[1]) * cos_angle

    # 椭圆的长短轴
    a = axes[0] / 2  # 长轴
    b = axes[1] / 2  # 短轴

    # 椭圆方程为 (x/a)^2 + (y/b)^2 = 1，求点 M 到椭圆的垂直距离
    distance_M_to_ellipse = (local_M_x ** 2 / a ** 2) + (local_M_y ** 2 / b ** 2) - 1
    if distance_M_to_ellipse < 0:
        distance_M_to_ellipse = -distance_M_to_ellipse  # 如果 M 在椭圆内，距离为负数，取绝对值
    p = np.sqrt(abs(distance_M_to_ellipse)) * min(a, b)  # 距离乘以椭圆的尺度因子

    return p


# 计算椭圆的前倾角
def calculate_antervision(major_axis_length, ellipse):
    center, axes, angle = ellipse
    # 找到点 M（大直径的五分之一处）
    D = major_axis_length
    M_x = center[0] + 0.2 * D * np.cos(np.radians(angle))
    M_y = center[1] + 0.2 * D * np.sin(np.radians(angle))
    M = (M_x, M_y)

    # 找到 M 到椭圆弧的垂直距离 p
    p = calculate_p(ellipse, M)

    # 计算前倾角 Antervision = arcsin(p / (0.4D))
    antervision_angle = np.arcsin(p / (0.4 * D))
    return antervision_angle  # 返回弧度值


# 保存角度信息到txt文件（中文输出）
def save_angles_to_file(json_path, abduction_angle_right, abduction_angle_left, antervision_right, antervision_left):
    # 构建结果保存路径
    output_folder = 'result/angle'
    os.makedirs(output_folder, exist_ok=True)

    # 创建文件名
    filename = os.path.basename(json_path).replace('.json', '.txt')
    output_path = os.path.join(output_folder, filename)

    # 写入文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(f"文件：{json_path}\n")

        # 外展角
        if abduction_angle_right is not None:
            f.write(f"右髋关节外展角：{abduction_angle_right:.4f} 弧度\n")
        else:
            f.write("右髋关节外展角：未计算\n")

        if abduction_angle_left is not None:
            f.write(f"左髋关节外展角：{abduction_angle_left:.4f} 弧度\n")
        else:
            f.write("左髋关节外展角：未计算\n")

        # 前倾角
        if antervision_right is not None:
            f.write(f"右髋关节前倾角：{antervision_right:.4f} 弧度\n")
        else:
            f.write("右髋关节前倾角：未计算\n")

        if antervision_left is not None:
            f.write(f"左髋关节前倾角：{antervision_left:.4f} 弧度\n")
        else:
            f.write("左髋关节前倾角：未计算\n")

    print(f"角度信息已保存到 {output_path}")


def process_json_file(json_path):
    data = load_json(json_path)

    # 提取关键点 1~6 作为右侧点，8~13 作为左侧点
    right_points, left_points = extract_keypoints(data)

    if len(right_points) < 2 or len(left_points) < 2:
        print(f"Skipping {json_path} due to insufficient points for abduction angle calculation.")
        return

    # 计算前倾角（拟合椭圆并使用前倾角公式）
    right_major_axis, right_center, right_ellipse = fit_ellipse_and_get_major_axis(right_points, json_path, "右侧")
    left_major_axis, left_center, left_ellipse = fit_ellipse_and_get_major_axis(left_points, json_path, "左侧")

    antervision_right = calculate_antervision(right_major_axis, right_ellipse) if right_major_axis is not None else None
    antervision_left = calculate_antervision(left_major_axis, left_ellipse) if left_major_axis is not None else None

    # 检查是否计算出了前倾角，并打印文件名
    if antervision_right is not None or antervision_left is not None:
        print(f"前倾角已成功计算：文件 {json_path}")

    # 计算右侧髋关节的外展角（检查是否有足够的点）
    if len(right_points) >= 6:
        abduction_angle_right = calculate_angle(right_points[1] - right_points[0], right_points[5] - right_points[4])  # 使用1~2 和5~6
    elif len(right_points) >= 2:  # 如果不足6个点，用最靠近的两个点
        abduction_angle_right = calculate_angle(right_points[1] - right_points[0], right_points[-1] - right_points[-2])
    else:
        abduction_angle_right = None

    # 计算左侧髋关节的外展角（检查是否有足够的点）
    if len(left_points) >= 6:
        abduction_angle_left = calculate_angle(left_points[1] - left_points[0], left_points[5] - left_points[4])  # 使用8~9 和12~13
    elif len(left_points) >= 2:  # 如果不足6个点，用最靠近的两个点
        abduction_angle_left = calculate_angle(left_points[1] - left_points[0], left_points[-1] - left_points[-2])
    else:
        abduction_angle_left = None

    # 保存计算结果到文件
    save_angles_to_file(json_path, abduction_angle_right, abduction_angle_left, antervision_right, antervision_left)


# 遍历所有文件夹并处理每个 JSON 文件
def process_all_folders(data_folder):
    for root, dirs, files in os.walk(data_folder):
        for file in files:
            if file.endswith('.json'):
                json_path = os.path.join(root, file)
                process_json_file(json_path)


# 主函数
def main():
    data_folder = 'image/DataMining_Dataset'  # 输入数据集的根文件夹
    process_all_folders(data_folder)


if __name__ == '__main__':
    main()
